from switchyard.lib.userlib import *
import random
import time

PERMIT = 0      #是否允许
PROTOCOL = 1    #协议
srcNET = 2
dstNET = 3
srcPORT = 4
dstPORT = 5
rateLIMIT = 6   #标志有无ratelimit
curTIME = 7     #当前时间
maxTOKEN = 8    #最大令牌数
curTOKEN = 9    #当前令牌数
IMPAIR = 10     #标志有无impair

#读取文件
def read_file():
    rules_list = []
    with open("lab_7/firewall_rules.txt") as r:
        rules = r.readlines()
        for i in rules:
            rule_list = i.strip('\n').split()
            if len(rule_list) <= 1 or rule_list[0] == '#': #注释
                pass
            else:
                rule_dict = {}
                rule_dict[PERMIT] = rule_list[0]
                if rule_list[1] == "ip" or rule_list[1] == "icmp":
                    rule_dict[PROTOCOL] = rule_list[1] 
                    rule_dict[srcNET] = rule_list[3]
                    rule_dict[dstNET] = rule_list[5]
                else:
                    rule_dict[PROTOCOL] = rule_list[1]
                    rule_dict[srcNET] = rule_list[3]
                    rule_dict[srcPORT] = rule_list[5]
                    rule_dict[dstNET] = rule_list[7]
                    rule_dict[dstPORT] = rule_list[9]
                
                if "ratelimit" in rule_list:
                    rule_dict[rateLIMIT] = 1 #标记
                    rule_dict[curTIME] = time.time()
                    rule_dict[maxTOKEN] = int(rule_list[-1]) * 2 #2r
                    rule_dict[curTOKEN] = int(rule_list[-1]) * 2
                else:
                    rule_dict[rateLIMIT] = 0

                if "impair" in rule_list:
                    rule_dict[IMPAIR] = 1 #标记
                else:
                    rule_dict[IMPAIR] = 0
                rules_list.append(rule_dict)
    return rules_list

#判断ip是否符合
def match_ip(txt_ip, ip):
    if txt_ip == "any" or IPv4Address(ip) in IPv4Network(txt_ip):
        return True
    else:
        return False

#判断端口是否符合
def match_port(txt_port, port):
    if txt_port == "any" or txt_port == port:
        return True
    else:
        return False

#扣除令牌数大小是否符合
def match_token(pkt, rule):
    if rule[rateLIMIT] == 1:
        size = len(pkt) - len(pkt.get_header(Ethernet))
        if size <= rule[curTOKEN]:
            rule[curTOKEN] -= size
            return True
        else:
            return False         
    return True

#判断丢包率是否符合
def match_random(pkt, rule):
    if rule[IMPAIR] == 1:
        a = random.randint(0,10)
        b = random.randint(0,5)
        if a < b:
            return False
    return True

#判断包是否符合防火墙条件
def match_pkt(pkt, rules_list):
    if pkt.has_header(IPv4):
        for rule in rules_list:
            if rule[PROTOCOL] == "ip":
                if match_ip(rule[srcNET], str(pkt[IPv4].src)) and match_ip(rule[dstNET], str(pkt[IPv4].dst)):
                    if match_token(pkt, rule) and match_random(pkt, rule):
                        if rule[PERMIT] == "permit":
                            return True
                        else:
                            return False
            elif pkt.has_header(ICMP):
                if rule[PROTOCOL] == "icmp":
                    if match_ip(rule[srcNET], str(pkt[IPv4].src)) and match_ip(rule[dstNET], str(pkt[IPv4].dst)):
                        if match_token(pkt, rule) and match_random(pkt, rule):
                            if rule[PERMIT] == "permit":
                                return True
                            else:
                                return False
            elif pkt.has_header(TCP):
                if rule[PROTOCOL] == "tcp":
                    if match_ip(rule[srcNET], str(pkt[IPv4].src)) and match_ip(rule[dstNET], str(pkt[IPv4].dst)) and match_port(rule[srcPORT], str(pkt[TCP].src)) and match_port(rule[dstPORT], str(pkt[TCP].dst)): 
                        if match_token(pkt, rule) and match_random(pkt, rule):
                            if rule[PERMIT] == "permit":
                                return True
                            else:
                                return False
            elif pkt.has_header(UDP):
                if rule[PROTOCOL] == "udp":
                    if match_ip(rule[srcNET], str(pkt[IPv4].src)) and match_ip(rule[dstNET], str(pkt[IPv4].dst)) and match_port(rule[srcPORT], str(pkt[UDP].src)) and match_port(rule[dstPORT], str(pkt[UDP].dst)):
                        if match_token(pkt, rule) and match_random(pkt, rule):
                            if rule[PERMIT] == "permit":
                                return True
                            else:
                                return False
    return True

def main(net):
    # assumes that there are exactly 2 ports
    portnames = [ p.name for p in net.ports() ]
    portpair = dict(zip(portnames, portnames[::-1]))
    rules_list = read_file()

    while True:
        pkt = None
        try:
            timestamp,input_port,pkt = net.recv_packet(timeout=0.5)
        except NoPackets:
            pass
        except Shutdown:
            break
        
        for rule in rules_list:
            if rule[rateLIMIT] == 1: 
                cur_time = time.time()
                if cur_time - rule[curTIME] > 0.5: #每0.5s向桶中添加令牌
                    rule[curTIME] = cur_time
                    rule[curTOKEN] += int(rule[maxTOKEN]/4)
                    rule[curTOKEN] = min(rule[curTOKEN], rule[maxTOKEN])
        if pkt is not None:
            # This is logically where you'd include some  firewall
            # rule tests.  It currently just forwards the packet
            # out the other port, but depending on the firewall rules
            # the packet may be dropped or mutilated.

            if match_pkt(pkt, rules_list):
                net.send_packet(portpair[input_port], pkt)

            
    net.shutdown()
